from langchain_core.runnables import RunnableBranch, RunnablePassthrough

branch_chain = RunnableBranch(
    (lambda x: 10 < len(x["text"]) < 20, lambda x: {"type": "long", **x}),
    (lambda x: len(x["text"]) < 10, lambda x: {"type": "short", **x}),
    lambda x: {"type": "default", **x},
)

chain = RunnablePassthrough.assign(proceed=branch_chain)

# invoke_param = {"text": "hello"}
# print(chain.invoke(invoke_param))

invoke_param = {"text": "hello world!"}
print(chain.invoke(invoke_param))

invoke_param = {"text": "hello world!hello world!hello world!hello world!hello world!"}
print(chain.invoke(invoke_param))
